{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "feef0c24",
   "metadata": {},
   "source": [
    "## step1:根据模板攻击\\机器学习得到的结果构建方程组"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2be004e9",
   "metadata": {},
   "source": [
    "- 模板攻击版本"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0d5af9ce",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-09T10:19:19.310778Z",
     "start_time": "2024-03-09T10:19:01.042678Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1571, 3)\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "##  ====带模的多项式乘法====  ##\n",
    "def poly_multi(X,m):\n",
    "    n = X.shape[0]\n",
    "    D1 = np.zeros((n,n),dtype = int)  \n",
    "    D2 = np.zeros((n,n),dtype = int)\n",
    "    \n",
    "    for j in range(m+1):    ## 矩阵运算法    coeff_front\n",
    "        D1[m][j] = X[m-j]\n",
    "\n",
    "    X = X[::-1]\n",
    "    for j in range(m+1,n):   ## 矩阵运算法    生成coeff_behind\n",
    "        D2[m][j] = X[j-m-1]\n",
    "\n",
    "    s1_coeff = (D1-D2)[m]\n",
    "    return s1_coeff \n",
    "################################################\n",
    "## =====构造结果矩阵（模板攻击）==== ##\n",
    "##:: 根据模板攻击得到的结果矩阵构造方程组\n",
    "## input: predict_file; metadat_file; z_file; save_file_url\n",
    "## output: result_fun\n",
    "def construct_result_fun(predict_file_url,metadata_file_url,z_file_url):\n",
    "    predict_file = np.load(predict_file_url)\n",
    "    print(predict_file.shape)\n",
    "    result_fun = np.zeros([predict_file.shape[0],258],dtype = int)  ## 每一行对应一个方程，0-255是系数；256是z值；257表示第几组的方程\n",
    "## 得到对应的方程：填充系数、Z值、第几组方程\n",
    "    for i in range(predict_file.shape[0]):\n",
    "        which_file = predict_file[i][0]\n",
    "        which_trace = predict_file[i][1]\n",
    "        which_poi = predict_file[i][2]\n",
    "        attack_file = np.load(metadata_file_url.format(which_file))\n",
    "        z_value = np.load(z_file_url.format(which_file))\n",
    "        poi = attack_file['index'][which_trace]\n",
    "        Z=z_value['Z']\n",
    "        C = z_value['C'][which_trace]  # 得到对应的C值\n",
    "        \n",
    "        result_fun[i,-1] = poi   # 第几组：0,1,2,3\n",
    "        if len(bin((Z[which_trace,poi,which_poi]) & 0xff)) == 10: ## 如果z低八位的二进制数是 11** **** ，说明对应的cs1的值为负值\n",
    "                result_fun[i,256] = ((Z[which_trace,poi,which_poi])&0xff)-256\n",
    "        else:\n",
    "            result_fun[i,256] = ((Z[which_trace,poi,which_poi]) & 0xff)\n",
    "        result_fun[i,0:256] = poly_multi(C,which_poi)   # 系数\n",
    "\n",
    "#     np.save(r\"F:\\paper_dilithium_0115\\function_file\\result_fun.npy\",result_fun)  ## 保存结果方程\n",
    "    return result_fun   \n",
    "#########################################################################\n",
    "## 把方程切分成4个\n",
    "## input: 总的方程矩阵\n",
    "## output: 4组方程\n",
    "def partition_equation(result_fun,save_file_url):\n",
    "    result_fun0 = np.zeros([1,258])\n",
    "    result_fun1 = np.zeros([1,258])\n",
    "    result_fun2 = np.zeros([1,258])\n",
    "    result_fun3 = np.zeros([1,258])\n",
    "    for i in range(result_fun.shape[0]):\n",
    "        if result_fun[i,-1] == 0:\n",
    "            result_fun0 = np.vstack((result_fun0,result_fun[i,:258]))\n",
    "        elif result_fun[i,-1] == 1:\n",
    "            result_fun1 = np.vstack((result_fun1,result_fun[i,:258]))\n",
    "        elif result_fun[i,-1] == 2:\n",
    "            result_fun2 = np.vstack((result_fun2,result_fun[i,:258]))\n",
    "        elif result_fun[i,-1] == 3:\n",
    "            result_fun3 = np.vstack((result_fun3,result_fun[i,:258]))\n",
    "\n",
    "    result_fun_group0 = result_fun0[1:,:].astype(int)\n",
    "    del result_fun0\n",
    "    result_fun_group1 = result_fun1[1:,:].astype(int)\n",
    "    del result_fun1\n",
    "    result_fun_group2 = result_fun2[1:,:].astype(int)\n",
    "    del result_fun2\n",
    "    result_fun_group3 = result_fun3[1:,:].astype(int)\n",
    "    del result_fun3        \n",
    "    np.save(save_file_url + \"\\\\result_fun_group0.npy\",result_fun_group0)\n",
    "    np.save(save_file_url + \"\\\\result_fun_group1.npy\",result_fun_group1)\n",
    "    np.save(save_file_url + \"\\\\result_fun_group2.npy\",result_fun_group2)\n",
    "    np.save(save_file_url + \"\\\\result_fun_group3.npy\",result_fun_group3)\n",
    "############################################################################################################################\n",
    "####                                 A     P     I                                                                      ####\n",
    "############################################################################################################################\n",
    "\n",
    "metadata_file_url=r\"G:\\RISCV_MLDSA\\Test_Metadata\\metadata_part{}.npz\"\n",
    "z_file_url=r\"G:\\RISCV_MLDSA\\Unpack\\Test\\Z_unpack_part{}.npz\"\n",
    "\n",
    "predict_file_url = r\"G:\\RISCV_MLDSA\\submit_code\\result.npy\"\n",
    "\n",
    "## 2.构造方程\n",
    "result_fun = construct_result_fun(predict_file_url,metadata_file_url,z_file_url)  # 调用函数,构造方程\n",
    "\n",
    "## 3.切割方程\n",
    "save_file_url = r\"G:\\RISCV_MLDSA\\functions\"    # 输出的四组方程矩阵的存储位置\n",
    "partition_equation(result_fun,save_file_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9231111",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-03T12:06:40.277822Z",
     "start_time": "2024-03-03T12:06:40.270142Z"
    }
   },
   "source": [
    "- 检查四组方程文件的正确方程数量（可选）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "dde823b5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-23T10:08:21.934559Z",
     "start_time": "2024-05-23T10:08:21.911764Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fun0 accuracy： 0.9289099526066351     ( 392 , 422 )  error:  30\n",
      "fun1 accuracy： 0.9105691056910569     ( 336 , 369 )  error:  33\n",
      "fun2 accuracy： 0.9219143576826196     ( 366 , 397 )  error:  31\n",
      "fun3 accuracy： 0.9138381201044387     ( 350 , 383 )  error:  33\n",
      "function number: 1571   real function: 1444  accuracy: 0.9191597708465945  error:  127\n"
     ]
    }
   ],
   "source": [
    "## 检验函数： 检查方程中正确的方程数量\n",
    "import numpy as np\n",
    "def check_fun(S1_file,check_fun_file_url):\n",
    "    S1 = np.load(S1_file)\n",
    "    total_fun_num = 0\n",
    "    total_real_fun_num = 0\n",
    "    for i in range(4):\n",
    "        fun_file = np.load(check_fun_file_url.format(i))\n",
    "        total_fun_num += fun_file.shape[0]\n",
    "        \n",
    "        real_num = 0\n",
    "        for j in range(fun_file.shape[0]):\n",
    "            if np.sum(S1[0,i] * fun_file[j,:256]) == fun_file[j,256]:\n",
    "                real_num += 1\n",
    "        print(\"fun{0} accuracy：\".format(i),real_num/fun_file.shape[0],\"    (\",real_num,\",\",fun_file.shape[0],\")\",\" error: \",fun_file.shape[0]-real_num)\n",
    "        total_real_fun_num += real_num\n",
    "    print(\"function number:\",total_fun_num,\"  real function:\",total_real_fun_num,\" accuracy:\",total_real_fun_num/total_fun_num,\" error: \",total_fun_num-total_real_fun_num)\n",
    "############################################################################################################################\n",
    "####                                 A     P     I                                                                      ####\n",
    "############################################################################################################################\n",
    "S1_file = r\"G:\\RISCV_MLDSA\\s1_attack.npy\"\n",
    "check_fun_file_url = r\"G:\\RISCV_MLDSA\\functions\\result_fun_group{}.npy\"\n",
    "check_fun(S1_file,check_fun_file_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "687e7a2f",
   "metadata": {},
   "source": [
    "## Step 2:使用线性规划求解器求解四组超定方程，获得S1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4aaf403",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-23T10:08:51.677070Z",
     "start_time": "2024-05-23T10:08:39.663768Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "function file:  (422, 258)\n",
      "Number of variables = 678\n",
      "Number of constraints = 844\n",
      "Solving with Glop solver v9.8.3296\n",
      "Solve successfully !!!\n",
      "Solution:\n",
      "Objective value = 420.0099999999999\n",
      "-------------------------------------------\n",
      "function file:  (369, 258)\n",
      "Number of variables = 625\n",
      "Number of constraints = 738\n",
      "Solving with Glop solver v9.8.3296\n",
      "Solve successfully !!!\n",
      "Solution:\n",
      "Objective value = 366.25999999999993\n",
      "-------------------------------------------\n",
      "function file:  (397, 258)\n",
      "Number of variables = 653\n",
      "Number of constraints = 794\n",
      "Solving with Glop solver v9.8.3296\n",
      "Solve successfully !!!\n",
      "Solution:\n",
      "Objective value = 394.35999999999984\n",
      "-------------------------------------------\n",
      "function file:  (383, 258)\n",
      "Number of variables = 639\n",
      "Number of constraints = 766\n",
      "Solving with Glop solver v9.8.3296\n",
      "Solve successfully !!!\n",
      "Solution:\n",
      "Objective value = 379.71999999999986\n",
      "-------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import random\n",
    "from ortools.linear_solver import pywraplp\n",
    "\n",
    "## --------------------------------------------------------------------------\n",
    "\n",
    "## ====线性规划函数==== ##\n",
    "## 求解四组方程，得到S1\n",
    "## input: 四组方程文件路径fun_file_url； 预测的S1文件存储位置save_S1_url；  第几组fun_i\n",
    "\n",
    "def BigM(fun_file_url,save_S1_url,fun_i):\n",
    "    const_fun = np.load(fun_file_url.format(fun_i))\n",
    "    print(\"function file: \",const_fun.shape)\n",
    "    # Create the mip solver with the SCIP backend.\n",
    "    # GLOP\n",
    "    solver = pywraplp.Solver.CreateSolver(\"GLOP\")\n",
    "    ## 定义变量\n",
    "\n",
    "    x = [solver.IntVar(-2, 2, f'x_{i}') for i in range(1, 257)]  \n",
    "\n",
    "    y = [solver.IntVar(0, 1,f'y_{i}') for i in range(1,const_fun.shape[0]+1)]\n",
    "\n",
    "    print(\"Number of variables =\", solver.NumVariables())\n",
    "\n",
    "    # 添加约束条件        # X = [-2 0 1 -1 1]\n",
    "    constraint_coeff = const_fun\n",
    "    constraints = []\n",
    "    k = 100\n",
    "    for i in range(constraint_coeff.shape[0]):\n",
    "        linear_expression = - constraint_coeff[i,256]\n",
    "        for j in range(256):\n",
    "            linear_expression += constraint_coeff[i,j] * x[j] \n",
    "        constraints.append(linear_expression <= k * (1 - y[i]))\n",
    "        constraints.append(linear_expression >= -k * (1 - y[i]))\n",
    "    #     constraints.append(constraint_coeff[i,0] * x[0] + constraint_coeff[i,1] * x[1] + constraint_coeff[i,2] * x[2] + constraint_coeff[i,3] * x[3] + constraint_coeff[i,4] * x[4] - constraint_coeff[i,5] <= k * (1-y_{i}) )\n",
    "    #     constraints.append(constraint_coeff[i,0] * x[0] + constraint_coeff[i,1] * x[1] + constraint_coeff[i,2] * x[2] + constraint_coeff[i,3] * x[3] + constraint_coeff[i,4] * x[4] - constraint_coeff[i,5] >= -k * (1-y_{i}) )\n",
    "\n",
    "    for constraint in constraints:\n",
    "        solver.Add(constraint)\n",
    "    print(\"Number of constraints =\", solver.NumConstraints())\n",
    "\n",
    "    ## 设定目标函数\n",
    "    solver.Maximize(sum(y))\n",
    "    # 设置时间限制\n",
    "    solver.set_time_limit(30000)\n",
    "    # 启用输出\n",
    "    solver.EnableOutput()\n",
    "    # 设置相对MIP间隙\n",
    "    # solver.parameters.SetDoubleParam('relative_mip_gap', 0.58)\n",
    "    check=[]\n",
    "    print(f\"Solving with {solver.SolverVersion()}\")\n",
    "    status = solver.Solve()\n",
    "    predict_S1 = np.zeros(256,dtype = float)\n",
    "    if status == pywraplp.Solver.OPTIMAL:\n",
    "        print(\"Solve successfully !!!\")\n",
    "        print(\"Solution:\")\n",
    "        print(\"Objective value =\", solver.Objective().Value())\n",
    "        i = 0\n",
    "        for variable in x:\n",
    "#             print(f\"{variable.name()} =\", variable.solution_value())\n",
    "            predict_S1[i] = variable.solution_value()\n",
    "            i += 1\n",
    "        print(\"-------------------------------------------\")\n",
    "        for variable in y:\n",
    "            # print(f\"{variable.name()} =\", variable.solution_value())\n",
    "            check.append(variable.solution_value())\n",
    "    else:\n",
    "        print(\"The problem does not have an optimal solution.\")\n",
    "    np.save(save_S1_url.format(fun_i),np.round(predict_S1))\n",
    "    return check\n",
    "############################################################################################################################\n",
    "####                                 A     P     I                                                                      ####\n",
    "############################################################################################################################\n",
    "fun_file_url = r\"G:\\RISCV_MLDSA\\functions\\result_fun_group{}.npy\"\n",
    "save_S1_url = r\"G:\\RISCV_MLDSA\\functions\\predict_fun_group{}.npy\"\n",
    "check1=[[] for _ in range(4)]\n",
    "for i in range(4):\n",
    "    chk=BigM(fun_file_url,save_S1_url,i)\n",
    "    check1[i].extend(chk)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "002749d1",
   "metadata": {},
   "source": [
    "- 检验求解得到的四组S1的正确率"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d1020bf8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-23T10:09:01.791507Z",
     "start_time": "2024-05-23T10:09:01.784495Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "predict S1_0 accuracy： 1.0     real: 256\n",
      "predict S1_1 accuracy： 1.0     real: 256\n",
      "predict S1_2 accuracy： 1.0     real: 256\n",
      "predict S1_3 accuracy： 1.0     real: 256\n",
      "Total S1 accuracy: 1.0   ( 1024 , 1024 )\n"
     ]
    }
   ],
   "source": [
    "## 检验S1： 检查预测的S1正确率\n",
    "def check_S1(S1_file,predict_S1_file_url):\n",
    "    S1 = np.load(S1_file)\n",
    "    \n",
    "    total_real_s1_num = 0\n",
    "    for i in range(4):\n",
    "        pre_s1 = np.load(predict_S1_file_url.format(i))\n",
    "        real_S1 = S1[0,i]\n",
    "        \n",
    "        real_num = np.sum(pre_s1 == real_S1)\n",
    "        \n",
    "        print(\"predict S1_{0} accuracy：\".format(i),real_num/256,\"    real:\",real_num)\n",
    "        total_real_s1_num += real_num\n",
    "    print(\"Total S1 accuracy:\",total_real_s1_num/1024,\"  (\",total_real_s1_num,\",\",1024,\")\")\n",
    "############################################################################################################################\n",
    "####                                 A     P     I                                                                      ####\n",
    "############################################################################################################################\n",
    "S1_file = r\"G:\\RISCV_MLDSA\\s1_attack.npy\"\n",
    "predict_S1_file_url = r\"G:\\RISCV_MLDSA\\functions\\predict_fun_group{}.npy\"\n",
    "check_S1(S1_file,predict_S1_file_url)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
